#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
技术面概率分析模块
通过历史K线数据计算买入后20日的上涨概率及收益率统计
"""

import requests
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from typing import Dict, Optional
import time


class EastMoneyKlineAPI:
    """东方财富K线API封装"""
    
    def __init__(self):
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
            'Referer': 'https://quote.eastmoney.com/',
            'Accept': 'application/json, text/plain, */*',
        })
    
    def get_kline_data(self, stock_code: str, days: int = 800) -> Optional[pd.DataFrame]:
        """
        获取股票历史K线数据（前复权）
        
        Args:
            stock_code: 股票代码，如 '600519'
            days: 获取多少天的数据，默认800天（约3年）
        
        Returns:
            DataFrame: 包含日期、开盘、收盘、最高、最低、成交量、成交额等
        """
        url = "https://push2his.eastmoney.com/api/qt/stock/kline/get"
        
        # 判断市场
        if stock_code.startswith('6'):
            secid = f"1.{stock_code}"  # 沪市/科创板
        elif stock_code.startswith('0'):
            secid = f"0.{stock_code}"  # 深市
        elif stock_code.startswith('3'):
            secid = f"0.{stock_code}"  # 创业板
        elif stock_code.startswith('688'):
            secid = f"1.{stock_code}"  # 科创板
        else:
            secid = f"1.{stock_code}"
        
        # 计算日期范围
        end_date = datetime.now().strftime('%Y%m%d')
        start_date = (datetime.now() - timedelta(days=days+365)).strftime('%Y%m%d')
        
        params = {
            'secid': secid,
            'ut': 'fa5fd1943c7b386f172d6893dbfba10b',
            'fields1': 'f1,f2,f3,f4,f5,f6',
            'fields2': 'f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61',
            'klt': '101',  # 日K线
            'fqt': '1',    # 前复权
            'beg': start_date,
            'end': end_date,
            'lmt': '10000'
        }
        
        try:
            response = self.session.get(url, params=params, timeout=30)
            response.raise_for_status()
            
            data = response.json()
            
            if data.get('rc') == 0 and data.get('data') and data['data'].get('klines'):
                klines = data['data']['klines']
                
                # 解析K线数据
                df_data = []
                for kline in klines:
                    fields = kline.split(',')
                    if len(fields) >= 7:
                        df_data.append({
                            '日期': fields[0],
                            '开盘': float(fields[1]) if fields[1] != '-' else 0,
                            '收盘': float(fields[2]) if fields[2] != '-' else 0,
                            '最高': float(fields[3]) if fields[3] != '-' else 0,
                            '最低': float(fields[4]) if fields[4] != '-' else 0,
                            '成交量': int(fields[5]) if fields[5] != '-' else 0,
                            '成交额': float(fields[6]) if fields[6] != '-' else 0,
                            '振幅': float(fields[7]) if len(fields) > 7 and fields[7] != '-' else 0,
                            '涨跌幅': float(fields[8]) if len(fields) > 8 and fields[8] != '-' else 0,
                            '涨跌额': float(fields[9]) if len(fields) > 9 and fields[9] != '-' else 0,
                            '换手率': float(fields[10]) if len(fields) > 10 and fields[10] != '-' else 0,
                        })
                
                if df_data:
                    df = pd.DataFrame(df_data)
                    df['日期'] = pd.to_datetime(df['日期'])
                    df = df.sort_values('日期').reset_index(drop=True)
                    return df
                
            return None
            
        except Exception as e:
            print(f"  ✗ 获取K线数据失败 ({stock_code}): {e}")
            return None


class ProbabilityAnalyzer:
    """技术面概率分析器"""
    
    def __init__(self):
        self.api = EastMoneyKlineAPI()
    
    def analyze_stock(self, stock_code: str, stock_name: str) -> Optional[Dict]:
        """
        对单只股票进行全面概率分析
        
        Args:
            stock_code: 股票代码
            stock_name: 股票名称
        
        Returns:
            分析结果字典，包含各种概率指标
        """
        print(f"  → 分析 {stock_code} {stock_name}...")
        
        # 获取历史K线数据
        df = self.api.get_kline_data(stock_code, days=800)
        if df is None or len(df) < 250:
            print(f"    ✗ 数据不足（需至少250个交易日）")
            return None
        
        print(f"    ✓ 获取 {len(df)} 个交易日数据")
        
        # 计算技术指标
        df = self._calculate_technical_indicators(df)
        
        # 计算20日后收益率
        df['20日后收盘'] = df['收盘'].shift(-20)
        df['20日收益率'] = ((df['20日后收盘'] - df['收盘']) / df['收盘']) * 100
        df['20日上涨'] = df['20日收益率'] > 0
        
        # 移除最后20行（没有未来数据）
        df_valid = df[:-20].copy()
        
        if len(df_valid) < 200:
            print(f"    ✗ 有效数据不足")
            return None
        
        # 当前状态
        current = df.iloc[-1]
        current_price = current['收盘']
        current_ma20 = current['MA20']
        current_ma60 = current['MA60']
        
        # 判断当前技术桶位
        bucket = self._classify_technical_bucket(current_price, current_ma20, current_ma60)
        
        # 统计分析
        # 1. 全样本统计
        全样本上涨概率 = (df_valid['20日上涨'].sum() / len(df_valid)) * 100
        全样本均值 = df_valid['20日收益率'].mean()
        全样本中位数 = df_valid['20日收益率'].median()
        全样本标准差 = df_valid['20日收益率'].std()
        
        # 2. 近1年统计
        近1年数据 = df_valid.tail(250)
        近1年上涨概率 = (近1年数据['20日上涨'].sum() / len(近1年数据)) * 100 if len(近1年数据) > 0 else 0
        近1年均值 = 近1年数据['20日收益率'].mean()
        近1年中位数 = 近1年数据['20日收益率'].median()
        
        # 3. 按技术桶位分层统计
        bucket_stats = self._analyze_by_bucket(df_valid)
        
        # 当前桶位的历史表现
        当前桶位数据 = df_valid[df_valid['技术桶位'] == bucket]
        if len(当前桶位数据) > 0:
            当前桶位上涨概率 = (当前桶位数据['20日上涨'].sum() / len(当前桶位数据)) * 100
            当前桶位均值 = 当前桶位数据['20日收益率'].mean()
            当前桶位样本量 = len(当前桶位数据)
        else:
            当前桶位上涨概率 = 全样本上涨概率
            当前桶位均值 = 全样本均值
            当前桶位样本量 = 0
        
        # 组装结果
        result = {
            '股票代码': stock_code,
            '股票名称': stock_name,
            '当前价格': current_price,
            'MA20': current_ma20,
            'MA60': current_ma60,
            '技术桶位': bucket,
            '技术桶位描述': self._get_bucket_description(bucket),
            
            # 概率指标
            '预估1月上涨概率': 当前桶位上涨概率,
            '全样本20日上涨概率': 全样本上涨概率,
            '近1年20日上涨概率': 近1年上涨概率,
            
            # 收益率统计
            '20日收益率_均值': 全样本均值,
            '20日收益率_中位数': 全样本中位数,
            '20日收益率_标准差': 全样本标准差,
            
            # 近1年统计
            '近1年收益率_均值': 近1年均值,
            '近1年收益率_中位数': 近1年中位数,
            
            # 当前桶位统计
            '当前桶位收益率_均值': 当前桶位均值,
            '当前桶位样本量': 当前桶位样本量,
            
            # 桶位分层统计
            '桶位统计': bucket_stats,
            
            # 原始数据（用于进一步分析）
            '历史数据': df,
        }
        
        print(f"    ✓ 分析完成")
        return result
    
    def _calculate_technical_indicators(self, df: pd.DataFrame) -> pd.DataFrame:
        """计算技术指标"""
        # MA20
        df['MA20'] = df['收盘'].rolling(window=20, min_periods=1).mean()
        
        # MA60
        df['MA60'] = df['收盘'].rolling(window=60, min_periods=1).mean()
        
        # 价格相对MA20位置
        df['价格相对MA20'] = ((df['收盘'] - df['MA20']) / df['MA20']) * 100
        
        # 价格相对MA60位置
        df['价格相对MA60'] = ((df['收盘'] - df['MA60']) / df['MA60']) * 100
        
        # 技术桶位分类
        df['技术桶位'] = df.apply(
            lambda row: self._classify_technical_bucket(row['收盘'], row['MA20'], row['MA60']),
            axis=1
        )
        
        return df
    
    def _classify_technical_bucket(self, price: float, ma20: float, ma60: float) -> str:
        """
        技术桶位分类
        
        Returns:
            'AboveBoth': 价格在MA20和MA60之上
            'Between': 价格在MA20和MA60之间
            'BelowBoth': 价格在MA20和MA60之下
        """
        if price >= ma20 and price >= ma60:
            return 'AboveBoth'
        elif price < ma20 and price < ma60:
            return 'BelowBoth'
        else:
            return 'Between'
    
    def _get_bucket_description(self, bucket: str) -> str:
        """获取桶位描述"""
        descriptions = {
            'AboveBoth': '双上（价格>MA20>MA60）',
            'Between': '中间（MA20<价格<MA60 或相反）',
            'BelowBoth': '双下（价格<MA20<MA60）'
        }
        return descriptions.get(bucket, '未知')
    
    def _analyze_by_bucket(self, df: pd.DataFrame) -> Dict:
        """按技术桶位分层统计"""
        buckets = ['AboveBoth', 'Between', 'BelowBoth']
        stats = {}
        
        for bucket in buckets:
            bucket_data = df[df['技术桶位'] == bucket]
            
            if len(bucket_data) > 0:
                stats[bucket] = {
                    '样本量': len(bucket_data),
                    '上涨概率': (bucket_data['20日上涨'].sum() / len(bucket_data)) * 100,
                    '平均收益率': bucket_data['20日收益率'].mean(),
                    '中位数收益率': bucket_data['20日收益率'].median(),
                    '标准差': bucket_data['20日收益率'].std(),
                }
            else:
                stats[bucket] = {
                    '样本量': 0,
                    '上涨概率': 0,
                    '平均收益率': 0,
                    '中位数收益率': 0,
                    '标准差': 0,
                }
        
        return stats


def test_single_stock():
    """测试单个股票"""
    analyzer = ProbabilityAnalyzer()
    
    # 测试张江高科
    result = analyzer.analyze_stock('600895', '张江高科')
    
    if result:
        print("\n" + "="*80)
        print(f"股票: {result['股票代码']} {result['股票名称']}")
        print(f"当前价格: {result['当前价格']:.2f}")
        print(f"技术桶位: {result['技术桶位描述']}")
        print(f"\n概率分析:")
        print(f"  预估1月上涨概率: {result['预估1月上涨概率']:.2f}%")
        print(f"  全样本20日上涨概率: {result['全样本20日上涨概率']:.2f}%")
        print(f"  近1年20日上涨概率: {result['近1年20日上涨概率']:.2f}%")
        print(f"\n收益率统计:")
        print(f"  均值: {result['20日收益率_均值']:.2f}%")
        print(f"  中位数: {result['20日收益率_中位数']:.2f}%")
        print(f"  标准差: {result['20日收益率_标准差']:.2f}%")
        print("="*80)


if __name__ == '__main__':
    test_single_stock()
